import asyncio
import socket
import json
import time
import zmq.asyncio
import time
from typing import Dict
from .common import *
from multiprocessing import Manager

class AsyncDeviceServer:
    def __init__(self, shared_devices: Dict[str, dict] = None):
        self.manager = Manager() if shared_devices is None else None
        self.devices: Dict[str, dict] = shared_devices if shared_devices is not None else self.manager.dict()
        # self.devices: Dict[str, dict] = {}
        self.context = zmq.asyncio.Context()
        self.sub_socket = self.context.socket(zmq.SUB)
        self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, "")
        self.connected = set()

    async def udp_listener(self):
        loop = asyncio.get_event_loop()
        udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        udp_socket.setblocking(False)
        udp_socket.bind(("", BROADCAST_PORT))

        while True:
            data, addr = await loop.sock_recvfrom(udp_socket, 1024)
            try:
                info = json.loads(data.decode())
                uuid_ = info['uuid']
                info['ip'] = addr[0]
                info['last_seen'] = time.time()
                info['status'] = 'online'
                # 只打印一次发现的设备
                if uuid_ not in self.devices:
                    self.devices[uuid_] = info
                    print(f"[UDP] 发现设备 {info['name']} @ {addr[0]}")
                else:
                    self.devices[uuid_].update(info)
            except Exception as e:
                print("[UDP] 数据解析失败：", e)

    async def zmq_subscriber(self):
        retry_counts = {}
        while True:
            for uuid_, info in self.devices.items():
                addr = f"tcp://{info['ip']}:{info['pub_port']}"
                if addr not in self.connected:
                    try:
                        if uuid_ not in retry_counts:
                            retry_counts[uuid_] = 0
                        if retry_counts[uuid_] < MAX_RETRY_COUNT:
                            self.sub_socket.connect(addr)
                            self.connected.add(addr)
                            print(f"[ZMQ] 订阅设备 {info['name']} 的心跳")
                            retry_counts[uuid_] = 0  # 重置重试计数
                        else:
                            print(f"[警告] 设备 {info['name']} 连接失败次数过多，等待下次发现")
                            retry_counts[uuid_] = 0  # 重置重试计数
                    except zmq.ZMQError as e:
                        retry_counts[uuid_] += 1
                        print(f"[ZMQ] 连接设备 {info['name']} 失败: {e}")
                        await asyncio.sleep(RECONNECT_DELAY)
            await asyncio.sleep(HEARTBEAT_INTERVAL)

            try:
                while True:
                    msg = await self.sub_socket.recv_string(flags=zmq.NOBLOCK)
                    heartbeat = json.loads(msg)
                    uuid_ = heartbeat['uuid']
                    if uuid_ in self.devices:
                        self.devices[uuid_]['last_seen'] = time.time()
                        if self.devices[uuid_]['status'] != 'online':
                            print(f"[恢复] 设备 {self.devices[uuid_]['name']} 恢复在线")
                        self.devices[uuid_]['status'] = 'online'
                        retry_counts[uuid_] = 0  # 重置重试计数
            except zmq.Again:
                await asyncio.sleep(HEARTBEAT_INTERVAL)

    async def device_status_checker(self):
        status_change_time = {}  # 记录状态变化时间
        offline_count = {}  # 记录连续离线次数
        while True:
            try:
                now = time.time()
                for uuid_, info in list(self.devices.items()):
                    time_diff = now - info.get('last_seen', 0)
                    current_status = info['status']
                    
                    # 检查是否需要更新状态
                    new_status = current_status
                    if time_diff > OFFLINE_TIMEOUT:
                        if uuid_ not in offline_count:
                            offline_count[uuid_] = 0
                        offline_count[uuid_] += 1
                        
                        # 只有连续多次检测到离线才更改状态
                        if offline_count[uuid_] >= 3:
                            new_status = 'offline'
                    else:
                        offline_count[uuid_] = 0
                        new_status = 'online'
                    
                    # 如果状态需要改变
                    if new_status != current_status:
                        # 检查状态是否在短时间内频繁变化
                        last_change = status_change_time.get(uuid_, 0)
                        if now - last_change > RECONNECT_DELAY:
                            info['status'] = new_status
                            status_change_time[uuid_] = now
                            
                            if new_status == 'offline':
                                print(f"[警告] 设备 {info['name']} 掉线")
                                # 从已连接集合中移除
                                addr = f"tcp://{info['ip']}:{info['pub_port']}"
                                if addr in self.connected:
                                    try:
                                        self.sub_socket.disconnect(addr)
                                        self.connected.remove(addr)
                                        print(f"[ZMQ] 断开与设备 {info['name']} 的连接")
                                    except zmq.ZMQError as e:
                                        print(f"[ZMQ] 断开连接失败: {e}")
                            else:
                                print(f"[恢复] 设备 {info['name']} 重新上线")
            except Exception as e:
                print(f"[错误] 设备状态检查异常: {e}")
            await asyncio.sleep(HEARTBEAT_INTERVAL)

    def get_online_devices(self):
        now = time.time()
        return [
            {
                "uuid": d["uuid"],
                "name": d["name"],
                "ip": d["ip"],
                "info": d.get("info", "")
            }
            for d in self.devices.values()
            if now - d.get("last_seen", 0) < 10
        ]

    async def start(self):
        print("[Server] 启动设备监听服务")
        await asyncio.gather(
            self.udp_listener(),
            self.zmq_subscriber(),
            self.device_status_checker()
        )

if __name__ == "__main__":
    server = AsyncDeviceServer()
    asyncio.run(server.start())